import torch
import init
import one_hot
from PIL import Image
from torchvision import transforms

def test_pic(path):
    img = Image.open(path)
    trans = transforms.Compose([
                transforms.ToTensor(),
                transforms.Resize((60, 160)),
                transforms.Grayscale()
            ])
    img_tensor = trans(img).cuda()
    img_tensor = img_tensor.reshape(1, 1, 60, 160)
    _model = torch.load("model.pth").cuda()
    output = _model(img_tensor)
    output = output.view(-1, init.captcha_array.__len__())
    output_label = one_hot.vec2Text(output)
    return output_label

if __name__ == '__main__':
    path = "./datasets/test/0wtn_1683891162.png"
    print(test_pic(path))